import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.nn.utils import prune

# 车速预测模型
class CarSpeedNet(nn.Module):
    def __init__(self):
        super(CarSpeedNet, self).__init__()
        self.fc1 = nn.Linear(4, 64)  # 输入为 [x, y, z, vr]，输出为 64 维向量
        self.fc2 = nn.Linear(64, 128)
        self.fc3 = nn.Linear(128, 3)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.3)

    def forward(self, x):
        x0=x
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.dropout(self.relu(self.fc2(x)))
        x = self.dropout(self.relu(self.fc3(x)))
        x = torch.mean(x,dim=1)#求出自车车速 base_link方向，相对local
        return x

    def _init_weights(self) -> None:
        # 嵌套的基本模块也会初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


# 蛙眼模型
class FrogEyeNet(nn.Module):
    def __init__(self, pre_model=None):
        super(FrogEyeNet, self).__init__()
        self.car_speed_net = pre_model

        self.conv0 = nn.Conv1d(4, 16,1)
        self.bn0 = nn.BatchNorm1d(16)
        self.conv1 = nn.Conv1d(16,64,1)
        self.bn1 = nn.BatchNorm1d(64)
        self.conv2 = nn.Conv1d(64,256,1)
        self.bn2 = nn.BatchNorm1d(256)

        # self.conv1 = nn.Conv1d(7, 32,1)
        # self.bn1 = nn.BatchNorm1d(32)
        # self.conv2 = nn.Conv1d(32, 64,1)
        # self.bn2 = nn.BatchNorm1d(64)
        self.conv3 = nn.Conv1d(260, 64,1)# 输出为动态/静态标记，可以视为二分类问题
        self.bn3 = nn.BatchNorm1d(64)
        if self.car_speed_net is None:
            self.conv4 = nn.Conv1d(64, 2,1)# 不使用车速信息
        else:
            self.conv4 = nn.Conv1d(67, 2,1)# 使用车速信息
        self.bn4 = nn.BatchNorm1d(2)
        self.relu = nn.ReLU()

    def forward(self, x):
        x0 = x
        if self.car_speed_net is None:
            xv_car0 = None
        else:
            xv_car0 = self.car_speed_net(x0)  # [b,n,4] -> [b,3]
            xv_car = xv_car0.unsqueeze(1).repeat(1,x0.shape[1],1).transpose(2,1)  # [b,3] -> [b,n,3]

        x = x.permute(0,2,1)
        x = self.relu(self.bn0(self.conv0(x)))
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x, _ = torch.max(x, dim=2)
        x = x.unsqueeze(1).repeat(1,100,1)
        x = torch.concat((x0, x), dim=2) # 点云自身特征加全局特征

        x= x.permute(0,2,1)
        x = self.relu(self.bn3(self.conv3(x)))
        if self.car_speed_net is None:
            pass
        else:
            x = torch.concat((xv_car, x), dim=1) # 特征加车速信息

        x = self.relu(self.bn4(self.conv4(x)))
        x = x.permute(0,2,1)
        return x, xv_car0

    def _init_weights(self) -> None:
        # 嵌套的基本模块也会初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

# 蛙眼检测模型
class FrogEyeDetNet(nn.Module):
    def __init__(self, pre_model=None):
        super(FrogEyeDetNet, self).__init__()
        self.frogeye = pre_model

        # self.conv0 = nn.Conv1d(4, 16,1)
        # self.bn0 = nn.BatchNorm1d(16)
        # self.conv1 = nn.Conv1d(16,64,1)
        # self.bn1 = nn.BatchNorm1d(64)
        # self.conv2 = nn.Conv1d(64,256,1)
        # self.bn2 = nn.BatchNorm1d(256)
        #
        # # self.conv1 = nn.Conv1d(7, 32,1)
        # # self.bn1 = nn.BatchNorm1d(32)
        # # self.conv2 = nn.Conv1d(32, 64,1)
        # # self.bn2 = nn.BatchNorm1d(64)
        # self.conv3 = nn.Conv1d(260, 64,1)# 输出为动态/静态标记，可以视为二分类问题
        # self.bn3 = nn.BatchNorm1d(64)
        # if self.car_speed_net is None:
        #     self.conv4 = nn.Conv1d(64, 2,1)# 不使用车速信息
        # else:
        #     self.conv4 = nn.Conv1d(67, 2,1)# 使用车速信息
        # self.bn4 = nn.BatchNorm1d(2)
        # self.relu = nn.ReLU()

    def forward(self, x):
        p_state, xv_car = self.frogeye(x) #输出 点的动静状态和自车车速(如果用了车速模型)
        p_state = nn.functional.softmax(p_state, dim=2) #点的动静状态转为动静概率值
        x = torch.concat((x, p_state), dim=2) #点的x y z vr 静概率 动概率
        return x

    def _init_weights(self) -> None:
        # 嵌套的基本模块也会初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
